package gossip.v2.core;


import gossip.v2.GossipCore;
import gossip.v2.cluster.GossipNode;
import lombok.Setter;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 状态管理
 */
public class StateManager {
    private final GossipNode localNode;
    private final Map<String, byte[]> state = new ConcurrentHashMap<>();
    private final Map<String, Long> versions = new ConcurrentHashMap<>();
    @Setter
    private GossipCore core;
    @Setter
    private ConflictResolver conflictResolver = new DefaultConflictResolver();

    public StateManager(GossipNode localNode) {
        this.localNode = localNode;
    }

    public byte[] getValueBytes(String key) {
        return state.get(key);
    }

    public Map<String, byte[]> getState() {
        return new ConcurrentHashMap<>(state);
    }

    public Map<String, Long> getVersions() {
        return new ConcurrentHashMap<>(versions);
    }

    public DeltaState getDeltaState(Map<String, Long> remoteVersion) {
        Map<String, Long> deltaVersions = new HashMap<>();
        Map<String, byte[]> deltaValues = new HashMap<>();

        versions.forEach((key, localVer) -> {
            long remoteVer = remoteVersion.getOrDefault(key, 0L);
            if (localVer > remoteVer) {
                deltaVersions.put(key, localVer);
                deltaValues.put(key, state.get(key));
            }
        });

        return new DeltaState(deltaVersions, deltaValues);
    }

    public int getStateDifferenceSize(Map<String, Long> remoteVersion) {
        int diffCount = 0;
        for (Map.Entry<String, Long> entry : versions.entrySet()) {
            long remoteVer = remoteVersion.getOrDefault(entry.getKey(), 0L);
            if (entry.getValue() > remoteVer) {
                diffCount++;
            }
        }
        return diffCount;
    }

    public void applyFullState(Map<String, Long> incomingVersions,
                               Map<String, byte[]> incomingValues) {
        incomingValues.forEach((key, remoteValue) -> {
            long remoteVer = incomingVersions.getOrDefault(key, 0L);
            long localVer = versions.getOrDefault(key, 0L);

            byte[] resolvedValue = remoteValue;
            if (conflictResolver != null && localVer > 0) {
                resolvedValue = (byte[]) conflictResolver.resolve(
                        key,
                        state.get(key),
                        remoteValue,
                        localVer,
                        remoteVer
                );
            }

            state.put(key, resolvedValue);
            versions.put(key, Math.max(localVer, remoteVer));
            notifyStateChange(key, resolvedValue);
        });
    }

    public Map<String, byte[]> mergeState(Map<String, Long> incomingVersions,
                                          Map<String, byte[]> incomingValues) {

        Map<String, byte[]> reverseDelta = new HashMap<>();

        incomingValues.forEach((key, remoteBytes) -> {
            long incomingVer = incomingVersions.getOrDefault(key, 0L);
            long localVer = versions.getOrDefault(key, 0L);

            if (incomingVer > localVer) {
                byte[] resolvedBytes = (byte[]) conflictResolver.resolve(
                        key,
                        state.get(key),
                        remoteBytes,
                        localVer,
                        incomingVer
                );

                state.put(key, resolvedBytes);
                versions.put(key, incomingVer);
                localNode.updateVersion(key, incomingVer);
                notifyStateChange(key, resolvedBytes);

            } else if (incomingVer < localVer) {
                reverseDelta.put(key, state.get(key));
            }
        });

        return reverseDelta;
    }

    public void update(String key, byte[] value) {
        long newVersion = versions.getOrDefault(key, 0L) + 1;
        state.put(key, value);
        versions.put(key, newVersion);
        localNode.updateVersion(key, newVersion);
        notifyStateChange(key, value);
    }

    private void notifyStateChange(String key, byte[] newValue) {
        core.getListeners().forEach(listener ->
                listener.onStateChanged(key, newValue));
    }

    public static class DeltaState {
        private final Map<String, Long> versions;
        private final Map<String, byte[]> stateValues;

        public DeltaState(Map<String, Long> versions, Map<String, byte[]> stateValues) {
            this.versions = versions;
            this.stateValues = stateValues;
        }

        public boolean isEmpty() {
            return stateValues.isEmpty();
        }

        public Map<String, Long> getVersions() { return versions; }
        public Map<String, byte[]> getStateValues() { return stateValues; }
    }

    @FunctionalInterface
    public interface ConflictResolver {
        Object resolve(String key, Object local, Object remote, long localVer, long remoteVer);
    }

    private static class DefaultConflictResolver implements ConflictResolver {
        @Override
        public Object resolve(String key, Object local, Object remote, long localVer, long remoteVer) {
            return remoteVer >= localVer ? remote : local;
        }
    }
}
